import numpy as np
import xarray as xr
from typing import Union, Tuple, Literal

def downsample(v: np.ndarray | xr.DataArray, interval:int,
               dims: int | str | list[int] | list[str]=None):
    """对 n 维数据进行降采样
    Downsample n-dimensional array

    Parameters
    ----------
    v : np.ndarray | xr.DataArray
        要降采样的数据
    interval : int
        降采样的间隔，每隔 interval 取一个值
    dims : int | list[int] | list[str]
        要降采样的维度，默认全部维度。可以是一个整数或者一组整数，如果数据是 DataArray，
        还可以是维度的名称
    """
    if isinstance(dims, (list, tuple)):
        if (not all(isinstance(i, int) for i in dims)) and \
                (not all(isinstance(i, str) for i in dims)):
            raise ValueError("dims must be list of int or str")
    
    if dims is None:
        slices = tuple(slice(None, None, interval) for _ in range(v.ndim))
        return v[slices]
    elif isinstance(dims, int):
        slices = tuple(slice(None, None, interval) if i == dims else slice(None)
                        for i in range(v.ndim))
        return v[slices]
    elif isinstance(dims, (list, tuple)):
        if isinstance(dims[0], int):
            slices = tuple(slice(None, None, interval) if i in dims else slice(None)
                            for i in range(v.ndim))
            return v[slices]
        elif isinstance(dims[0], str):
            if isinstance(v, xr.DataArray):
                slices = {i: slice(None, None, interval) for i in dims}
                return v.isel(slices)
            else:
                raise ValueError("only DataArray support str dims")
    else:
        raise ValueError("dims must be int, list of int or str")

Array = Union[np.ndarray, xr.DataArray, xr.Dataset]
FilterType = Literal['lowpass', 'highpass', 'bandpass', 'bandstop']

def _butter(data: np.ndarray, T: float|Tuple[float, float], fs: int, 
            btype: FilterType, axis: int = 0, N: int = 3,) -> np.ndarray:
    """巴特沃斯滤波器时间周期转换"""
    import scipy.signal as signal

    if btype in ['lowpass', 'highpass']:
        Wn = 2/(T*fs)
    elif btype in ['bandpass', 'bandstop']:
        Wn = (2/(T[1]*fs), 2/(T[0]*fs))
    else:
        raise ValueError("ftype参数错误, "
                "有效值为 'lowpass', 'highpass', 'bandpass', 'bandstop'")

    b, a = signal.butter(N, Wn, btype)
    return signal.filtfilt(b, a, data, axis=axis)


def _get_axis(data: xr.DataArray, axis: int | str) -> int:
    """获取数据某个维度的索引，可以是维度的索引值，也可以是维度的名称"""
    if isinstance(axis, str):
        axis = data.dims.index(axis)
    return axis


def butter_filter(data: Array, T: float | Tuple[float, float], 
                  fs: int, btype: FilterType, axis: int | str = 0, N: int = 3,) -> Array:
    """巴斯沃斯滤波，输入的周期与频率使用相同的时间单位。

    例如要对数据进行 5-10 日带通滤波，选取时间单位为日，那么参数 T=(5, 10)，
    数据的时间分辨率是 2 小时，那么 fs=12 （一个时间单位内的采样点数）。


    Butterworth filter, the input period and frequency use the same time unit.

    For example, to perform a 5-10 day bandpass filter on the data,
    set T=(5, 10) with time unit in days, and if the data
    has a time resolution of 2 hours, then fs=12 (number of samples per time unit).

    Reference: https://blog.csdn.net/weixin_44237337/article/details/124135064

    Parameters
    ----------
    data : Array
        需要滤波的数据，可以是 numpy.ndarray、xarray.DataArray 或 xarray.Dataset
        其中 xarray.Dataset 中的每个变量都会被滤波。

        The data to be filtered, can be numpy.ndarray, xarray.DataArray or 
        xarray.Dataset.
        Note that each variable in xarray.Dataset will be filtered.

    T : float | Tuple[float, float]
        滤波的周期，低通、高通为单个浮点数或整数，带通、带阻为一个元组。
        例如 T=(5, 10) 表示 5-10 日带通，注意：T 的单位与 fs 的单位相同。

        The period of the filter, a single float or int for lowpass/highpass,
        or a tuple for bandpass/bandstop.
        For example, T=(5, 10) means a bandpass filter from 5 to 10 days.
        Note that the unit of T must match the unit of fs.

    fs : int
        每个单位时间内的采样点数，例如数据的时间分辨率是 2 小时，滤波周期以日为单位，
        则一日中有 12 次采样，fs=12。

        The number of samples per time unit, for example, if the data has a time
        resolution of 2 hours and the filter period is in days, then there are
        12 samples in a day, so fs=12.

    btype : FilterType
        滤波类型，'lowpass', 'highpass', 'bandpass', 'bandstop'

        The type of filter, can be 'lowpass', 'highpass', 'bandpass', or 
        'bandstop'.

    axis : int | str, optional
        滤波的轴，可以是轴的索引值或轴的名称，默认为 0。

        The axis to filter, can be an index or a name, default is 0.

    N : int, optional
        滤波器的阶数，默认为 3。

        The order of the filter, default is 3.

    Returns
    -------
    Array
        返回与输入数据类型相同的滤波后的数据。

        Returns the filtered data with the same type as the input data.
    """
    
    if isinstance(data, np.ndarray):
        return _butter(data, T, fs, btype, axis, N)

    elif isinstance(data, xr.DataArray):
        axis = _get_axis(data, axis)
        data.values = _butter(data.values, T, fs, btype, axis, N)
        return data
    
    elif isinstance(data, xr.Dataset):
        # 对数据集的每一个变量进行滤波
        for key in data.data_vars.keys():
            axis = _get_axis(data[key], axis)
            data[key].values = _butter(data[key].values, T, fs, btype, axis, N)
        return data
